from mcp.server.fastmcp import FastMCP
import pandas as pd
from typing import List, Dict
import matplotlib.pyplot as plt
import numpy as np

# Create an MCP server
mcp = FastMCP("Stock Prediction Server")

@mcp.tool()
def predict_stock(data: List[Dict], model_type: str = "ARIMA", days: int = 7) -> Dict:
    """
    基于历史数据预测未来股市价格
    Args:
        data: 历史数据列表或CSV文件路径(可带#行号标记)，包含日期、收盘价等字段
        model_type: 预测模型类型 (ARIMA或LSTM)
        days: 预测天数
    Returns:
        包含预测结果的字典
    """
    if isinstance(data, str):
        # 处理CSV格式数据或文件路径
        import os
        from io import StringIO
        
        if '#' in data and os.path.exists(data.split('#')[0]):
            # 处理带行号标记的文件路径输入
            file_path = data.split('#')[0]
            line_range = data.split('#')[1] if '#' in data else None
            
            # 读取文件并处理行号范围
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                
                if line_range and line_range.startswith('L'):
                    start_line, end_line = map(int, line_range[1:].split('-'))
                    csv_content = ''.join(lines[start_line-1:end_line])
                    df = pd.read_csv(StringIO(csv_content))
                else:
                    df = pd.read_csv(file_path)
            
            # 将DataFrame转换为字典列表格式
            data = df.to_dict('records')
        elif os.path.exists(data):
            # 处理普通文件路径输入
            df = pd.read_csv(data)
            data = df.to_dict('records')
        else:
            # 处理CSV字符串输入
            try:
                # 尝试解析为JSON列表
                import json
                data_list = json.loads(data)
                if isinstance(data_list, list):
                    df = pd.DataFrame(data_list)
                    data = df.to_dict('records')
                else:
                    df = pd.read_csv(StringIO(data))
                    data = df.to_dict('records')
            except json.JSONDecodeError:
                df = pd.read_csv(StringIO(data))
                data = df.to_dict('records')
    else:
        df = pd.DataFrame(data)
        data = df.to_dict('records')
    # 这里实现预测逻辑
    predictions = {"dates": [], "prices": []}
    
    # 示例预测结果
    last_price = df.iloc[-1]["收盘"]
    from datetime import datetime, timedelta
    # 转换日期格式从YYYYMMDD到YYYY-MM-DD
    last_date_str = df.iloc[-1]["日期"]
    last_date = datetime.strptime(f"{last_date_str[:4]}-{last_date_str[4:6]}-{last_date_str[6:]}", '%Y-%m-%d')
    for i in range(1, days+1):
        pred_date = last_date + timedelta(days=i)
        predictions["dates"].append(pred_date.strftime('%Y-%m-%d'))
        predictions["prices"].append(last_price * (1 + np.random.uniform(-0.02, 0.03)))
    
    # 生成可视化图表
    plt.figure(figsize=(10, 6))
    plt.plot(predictions["dates"], predictions["prices"], marker='o')
    plt.title(f"Stock Price Prediction ({model_type})")
    plt.xlabel("Date")
    plt.ylabel("Price")
    plt.grid(True)
    plt.xticks(rotation=45)
    
    # 保存图表为PNG文件
    plt.savefig('stock_prediction.png')
    plt.close()
    
    return {
        "model": model_type,
        "predictions": predictions,
        "visualization": "图表已保存为stock_prediction.png文件"
    }

@mcp.tool()
def risk_alert(data: List[Dict], threshold: float = 0.05) -> Dict:
    """
    基于波动率分析识别高风险时段
    Args:
        data: 历史数据列表
        threshold: 波动率阈值
    Returns:
        包含风险等级和预警信息的字典
    """
    df = pd.DataFrame(data)
    returns = df["涨跌幅"] / 100  # 假设涨跌幅是百分比
    volatility = returns.std()
    
    risk_level = "high" if volatility > threshold else "medium" if volatility > threshold/2 else "low"
    
    return {
        "volatility": volatility,
        "risk_level": risk_level,
        "alert": "High risk" if risk_level == "high" else "Normal",
        "suggestion": "Consider reducing position" if risk_level == "high" else "No action needed"
    }

@mcp.resource("stock://{symbol}/history")
def get_stock_history(symbol: str) -> List[Dict]:
    """
    获取股票历史数据资源
    Args:
        symbol: 股票代码
    Returns:
        历史数据列表
    """
    # 实际实现中应从数据库或API获取数据
    return []

if __name__ == "__main__":
    mcp.run(transport="stdio")